import gradio as gr
import numpy as np
import pickle
import cv2
import joblib  # 用于加载joblib保存的模型

# 加载最佳模型
with open('D:/tuxiang/faiss_dog_cat_question_finish/best_model.pkl', 'rb') as f:
    model = joblib.load(f)

# 定义预测函数
def predict_image(image):
    # 如果是彩色图片，则转为灰度图
    if len(image.shape) == 3: 
        img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:  
        img = image
    
    # 调整图像尺寸为32x32，与之前的模型输入一致
    img = cv2.resize(img, (32, 32))
    
    # 将图像扁平化处理，以适应模型输入格式
    image_flattened = img.flatten().reshape(1, -1)
    
    # 使用模型进行预测
    prediction = model.predict(image_flattened)[0]
    
    # 返回预测结果
    return {
        "猫": float(prediction == 0),
        "狗": float(prediction == 1)
    }

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_image,  # 使用自定义的预测函数
    inputs=gr.Image(),  # 输入为图片
    outputs=gr.Label(num_top_classes=2)  # 输出标签，显示前两个类的概率
)

# 启动应用
iface.launch()
